from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_experimental.utilities import PythonREPL
from langchain_openai import ChatOpenAI

chat_model = ChatOpenAI(
    openai_api_key="key",
    openai_api_base="https://api.moonshot.cn/v1",
    model="moonshot-v1-8k",
    temperature=0,
    request_timeout=60,
    max_retries=3,
)
template = """编写一些Python代码来解决用户的问题。

只返回Markdown格式的Python代码，例如：

```python
....
```"""
prompt=ChatPromptTemplate.from_messages([
    ("system",template),("human","{input}")
])

def _sanitize_output(text:str):
    _,after=text.split("```python")
    # print(text)
    # print(after.split("```")[0])
    return after.split("```")[0]

chain = prompt | chat_model | StrOutputParser() | _sanitize_output | PythonREPL().run
print(chain.invoke({"input": "2加2等于多少"}))